package com.yanggu.spark.sql

import org.apache.spark.sql.Row
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types._

/**
 * 实现方式 - UDAF - 弱类型
 * 定义类继承UserDefinedAggregateFunction，并重写其中方法
 *
 */
class MyAverageUDAF extends UserDefinedAggregateFunction {

  /**
   * 聚合函数输入参数的数据类型
   *
   * @return
   */
  override def inputSchema: StructType = StructType(Array(StructField("age", IntegerType)))

  /**
   * 聚合函数缓冲区中值的数据类型(age,count)
   *
   * @return
   */
  override def bufferSchema: StructType = StructType(Array(StructField("sum", LongType), StructField("count", LongType)))

  /**
   * 返回的数据类型
   *
   * @return
   */
  override def dataType: DataType = DoubleType

  /**
   * 稳定性：对于相同的输入是否一直返回相同的输出。
   * @return
   */
  override def deterministic: Boolean = true

  /**
   * 函数缓冲区初始化
   * @param buffer
   */
  override def initialize(buffer: MutableAggregationBuffer): Unit = {
    // 存年龄的总和
    buffer(0) = 0L
    // 存年龄的个数
    buffer(1) = 0L
  }

  /**
   * 当数据输入时如何处理
   * @param buffer 缓冲区里面的数据
   * @param input 输入的数据
   */
  override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
    if (!input.isNullAt(0)) {
      buffer(0) = buffer.getLong(0) + input.getInt(0)
      buffer(1) = buffer.getLong(1) + 1
    }
  }

  /**
   * 合并缓冲区
   * @param buffer1
   * @param buffer2
   */
  override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
    buffer1(0) = buffer1.getLong(0) + buffer2.getLong(0)
    buffer1(1) = buffer1.getLong(1) + buffer2.getLong(1)
  }

  /**
   * 计算最终的结果
   * @param buffer
   * @return
   */
  override def evaluate(buffer: Row): Any = {
    buffer.getLong(0).toDouble / buffer.getLong(1)
  }
}
